mod attr_retry;
use attr_retry::{expand_retry, Args};

mod utils;
#[allow(unused_imports)]
use utils::print_log;

use anyhow::{bail, Result};
use proc_macro::TokenStream;
use proc_macro2::TokenTree;
use quote::quote;
use syn::{parse_macro_input, DeriveInput, Expr, FnArg, Ident, ItemFn};

#[proc_macro_derive(Hello)]
/// impl struct Hello with a hello function
pub fn hello_derive(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;
    let expanded = quote! {
        impl #name {
            pub fn hello() {
                println!("Hello, world!");
            }
        }
    };
    expanded.into()
}

#[proc_macro_derive(AnswerFn)]
/// 增加一个answer函数
pub fn derive_answer_fn(_input: TokenStream) -> TokenStream {
    "fn answer() -> i32 { 42 }".parse().unwrap()
}

#[proc_macro_derive(PrintLog)]
/// 打印日志
pub fn derive_print_log(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = &input.ident;
    let vis = &input.vis;
    let attrs = &input.attrs;

    // 格式化输出
    let attrs_formatted = format!("{:?}", attrs);
    let vis_formatted = format!("{:?}", vis);
    let data_formatted = format!("{:?}", input.data);
    let expanded = quote! {
        impl #name {
            /// 打印name, data, vis
            pub fn print_log() {
                println!("name: {}",stringify!(#name));
                println!("data: {:?}",#data_formatted);
                println!("vis: {}",#vis_formatted);
            }
            /// 打印attrs, 但是是个空[]
            pub fn print_attrs() {
                    println!("attr: {:?}",#attrs_formatted);
            }
        }
    };
    expanded.into()
}

#[proc_macro_attribute]
/// 打印日志
pub fn retry(attr: TokenStream, item: TokenStream) -> TokenStream {
    let item_fn = parse_macro_input!(item as ItemFn);
    let args = parse_macro_input!(attr as Args);
    let token_stream = expand_retry(args.attr, item_fn);
    token_stream
        .unwrap_or_else(|e| e.into_compile_error())
        .into()
}

/// # 属性式过程宏
///
/// ## `pub struct TokenStream { /* private fields */ }`
///
/// An abstract stream of tokens, or a sequence of token trees.
///
/// This type provides an interface for iterating over token trees, and collecting token trees into a stream.
///
/// TokenStream is both the input and output of `#[proc_macro]`, `#[proc_macro_attribute]` and `#[proc_macro_derive]` definitions.
///
/// ## `pub fn parse<T: Parse>(input: TokenStream) -> Result<T>`
///
/// Parse tokens of source code into the chosen syntax tree node.
///
/// ## have to parsing attr
///
/// 类似python装饰器的宏 `deco`
#[proc_macro_attribute]
pub fn deco(attr: TokenStream, item: TokenStream) -> TokenStream {
    // // 把item 转成 Tokens: TokenStream
    // let func = item.into();
    // // 把func(Tokens: TokenStream)转为ItemFn
    // let item_fn: ItemFn = syn::parse(func).expect("Input is not a function");
    // 上面2句与下面1句等效  get the ItemFn
    // ItemFn: A free-standing function definition.
    let item_fn = parse_macro_input!(item as ItemFn);
    // get the visibility of the function
    let vis = &item_fn.vis;
    // get the function name as Ident
    let ident = &item_fn.sig.ident;
    // ident 是被装饰的函数, 在例子中是add2函数
    // get the function body, Box<Block>
    let block = &item_fn.block;

    // 获取参数 Punctuated<FnArg, Comma>
    // the arguments of the decorated function
    // in this case, it is a single argument "a: i32"
    let inputs = item_fn.sig.inputs;
    // 获取返回类型 ReturnType
    // -> i32 {
    let output = item_fn.sig.output;

    // print_log(vis, ident, block, &inputs, &output);

    // 遍历参数inputs, 得到参数名
    //
    // ```rest
    // pub struct PatType {
    //     pub attrs: Vec<Attribute>,
    //     pub pat: Box<Pat>,
    //     pub colon_token: Token![:],
    //     pub ty: Box<Type>,
    // }
    // ```
    //
    // `syn::FnArg::Typed(PatType)`
    //
    // ```rust
    // pub enum Pat { 17 variants}  // such as Ident, Path, Lit, etc.
    // ```
    //
    // 被装饰的函数的参数
    // The arguments of the decorated function
    // in this case, it is a single argument "a: i32"
    //
    let input_values: Vec<_> = inputs
        .iter()
        .map(|arg| match arg {
            FnArg::Typed(ref val) => &val.pat,
            _ => unimplemented!("#[deco] can not be used with associated function"),
        })
        .collect();
    // println!("input_values: {:#?}", input_values);
    // println!("{}", "#".repeat(40));

    // 打印attr
    //
    // ```text
    // attr_test: TokenStream [
    //     Ident {
    //         ident: "logging",
    //         span: #0 bytes(267..274),
    //     },
    // ]
    // ```
    // println!("attr_test: {:#?}", attr);
    // println!("{}", "#".repeat(40));

    // custom attribute parsing
    // #[deco(logging)]
    //        ~~~~~~~
    // 'logging' is the attr. no args in this case.
    let attr = DecoratorAttr::parse(attr.into()).expect("Failed to parse attribute");
    let caller = match attr {
        DecoratorAttr::Fixed { name } => {
            quote::quote! {
                // ! The core part is rewriting the decorated function.
                // `#ident` represents the decorated function, which is "add2" in this case
                // equivalent to:
                // * `visibility inherent fn add2(a: i32) -> i32 {`
                #vis fn #ident(#inputs) #output {
                    // #name represents the decorator function, which is "logging" in this case
                    // - equivalent to:
                    // * `let f = logging(deco_internal)`;
                    let f = #name(deco_internal);
                    // 这个语法 #(...)* 将对 input_values 中的每个元素展开为函数调用参数
                    // - equivalent to:
                    // * `return f(a); // this means logging(add2(a))`
                    return f(#(#input_values,)*);

                    // deco_internal is the decorated function, which is "add2" in this case
                    // - equivalent to:
                    // * ```rust
                    // * fn deco_internal(a: i32) -> i32 {
                    // *     // add2 function body
                    // *     a + 2
                    // * }
                    // * ```
                    fn deco_internal(#inputs) #output #block
                }
            }
        }
        DecoratorAttr::Parametric { name, args } => {
            quote::quote! {
            // ```rust
            // #vis fn add2(a: i32) -> i32 {
            //     let deco = logging(args);  // there is no args here actually
            //     let f = deco(deco_internal);  // logging(args)(deco_internal)
            //     return f();
            //
            //     fn deco_internal(a: i32) -> i32 {
            //         a + 2
            //     }
            // }
            // ```
            #vis fn #ident(#inputs) #output {
                let deco = #name(#(#args,)*);
                let f = deco(deco_internal);
                return f(#(#input_values,)*);

                fn deco_internal(#inputs) #output #block
            }
            }
        }
    };
    caller.into()
}

#[derive(Debug, PartialEq)]
/// 装饰器参数的数据结构,
/// use enum data structure to represent two types of decorator attributes:
/// 一种是只有名字的固定装饰器，另一种是带参数的可变装饰器
enum DecoratorAttr {
    Fixed { name: Ident },
    Parametric { name: Ident, args: Vec<Expr> },
}

impl DecoratorAttr {
    /// 解析装饰器参数
    /// 注意! 这里用的是proc_macro2::TokenStream，而不是proc_macro::TokenStream
    fn parse(attr: proc_macro2::TokenStream) -> Result<Self> {
        // 空标识符
        let mut ident = None;
        // 空参数列表
        let mut args = Vec::new();
        // 遍历参数attr
        // TokenTree 是 TokenStream 的一个元素，可以是标识符、字面量、括号等
        //
        // ```rust
        // // TokenTree 是一个枚举
        // pub enum TokenTree {
        //     Group(Group),  // 包含分界符的token steam
        //     Ident(Ident),  // 标识符
        //     Punct(Punct),  // 单个的标点
        //     Literal(Literal),  // 字面量, 'a',"hello", 123, 3.14
        // }
        // ```
        for at in attr {
            match at {
                TokenTree::Ident(id) => {
                    ident = Some(id);
                }
                TokenTree::Group(grp) => {
                    if ident.is_none() {
                        bail!("Invalid token stream");
                    }
                    for t in grp.stream() {
                        // parse2 解析 proc_macro2::TokenTree 到 AST
                        if let Ok(expr) = syn::parse2(t.into()) {
                            // args里存表达式
                            args.push(expr);
                        }
                    }
                }
                _ => {
                    bail!("Invalid token stream");
                }
            }
        }
        // 返回DecoratorAttr的结果
        if let Some(name) = ident {
            if args.is_empty() {
                Ok(DecoratorAttr::Fixed { name })
            } else {
                Ok(DecoratorAttr::Parametric { name, args })
            }
        } else {
            bail!("Invalid token stream");
        }
    }
}

#[proc_macro_attribute]
/// Practice of using TokenStream to parse attribute and function
pub fn my_log(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let item_fn = parse_macro_input!(item as ItemFn);
    let _attr = parse_macro_input!(_attr as Args);

    let vis = &item_fn.vis;
    let ident = &item_fn.sig.ident;
    let block = &item_fn.block;
    let inputs = item_fn.sig.inputs;
    let output = item_fn.sig.output;

    // Parse the input_values from Punctuated<FnArg, Comma> to Vec<&Pat>
    let input_values: Vec<_> = inputs
        .iter()
        .map(|arg| match arg {
            FnArg::Typed(ref val) => &val.pat,
            _ => unimplemented!("#[my_log] can not be used with associated function"),
        })
        .collect();

    // ! Generate the output code
    let expanded = quote! {
        // fn add_two_numbers(a: i32, b: i32) -> i32 {
        //     return inside(11,22);
        //     fn inside(a: i32, b: i32) -> i32 {
        //         a + b
        //     }
        // }
        #vis fn #ident(#inputs) #output{
            println!("Input: {}", stringify!(#(#input_values),*));
            let result = inside(#(#input_values, )*);
            println!("Output: {:?}", result);
            // 2 conditions: 1. no args, 2. with args
            return result;
            fn inside(#inputs) #output #block
        }
    };
    expanded.into()
}
